/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg remote update.
 * Author: qingjicheng
 * Create: 2022-12-17
 */

#ifndef REMOTE_UPDATE_H
#define REMOTE_UPDATE_H
#include <algorithm>
#include <atomic>
#include "alg_embedding_service.h"

namespace hccl {

class AlgCollRemoteUpdate : public AlgEmbeddingService {
public:
    explicit AlgCollRemoteUpdate();
    ~AlgCollRemoteUpdate() override;
    HcclResult Process(HcclOpPtr opPtr) override;

    HcclResult AllocTransferMemAndRegMrForPs(const HcclOperator<EmbeddingServiceParam> &op);

private:
    HcclResult AlgRemoteUpdate(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult AlgRecvUpdateRequest(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult AlgIsendUpdateResponse(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult AlgWaitSome(HcclOperator<EmbeddingServiceParam> *opPtr);
    HcclResult AlgUpdate(HcclOperator<EmbeddingServiceParam> *opPtr);

    FuncTable funcTable_;

    std::atomic<bool> transferMemAllocated_{};
    void *keyTransferMem_{};
    void *valueTransferMem_{};

    std::unique_ptr<s8[]> keyTransferBeginMem_{};
    std::unique_ptr<s8[]> valueTransferBeginMem_{};
};
}
#endif
